#!/bin/bash

BLKON=test/online
BLKOFF=test/offline

LVM_ON=/dev/mapper/docker-iocost_online
LVM_OFF=/dev/mapper/docker-iocost_offline

DIRON=/export/iocost_online
DIROFF=/export/iocost_offline

WEIGHT_ONLINE=1000
WEIGHT_OFFLINE=250


function disable_read_ahead()
{
  dev=$1
  if [ -n "${dev}" ]; then
    NUM=$(lsblk -no MAJ:MIN ${dev})
    d=$(echo $NUM)
    echo 0 > /sys/class/bdi/${d}/read_ahead_kb
  else
    echo "no device specified, do nothing."
  fi
}

function run_fio()
{
  local rw=$1
  local direct=$2

  cgexec -g cpu,memory,blkio:${BLKON} fio -directory=${DIRON} -direct=${direct} -iodepth=1 -thread -rw=${rw} -ioengine=psync -bs=16k -size=1G -numjobs=16 -group_reporting -name=mytest_online &> online-${rw}.log &

  cgexec -g cpu,memory,blkio:${BLKOFF} fio -directory=${DIROFF} -direct=${direct} -iodepth=1 -thread -rw=${rw} -ioengine=psync -bs=16k -size=1G -numjobs=16 -group_reporting -name=mytest_offline &> offline-${rw}.log &
}

function fio_direct_seq_write()
{
  echo "run fio test (direct seq write)"
  run_fio write 1
}

function fio_direct_rand_write()
{
  echo "run fio test (direct rand write)"
  run_fio randwrite 1
}

function fio_direct_seq_read()
{
  echo "run fio test (direct seq read)"
  run_fio read 1
}

function fio_direct_rand_read()
{
  echo "run fio test (direct rand read)"
  run_fio randread 1
}

function fio_buffer_seq_write()
{
  echo "run fio test (buffer seq write)"
  run_fio write 0
}

function fio_buffer_rand_write()
{
  echo "run fio test (buffer rand write)"
  run_fio randwrite 0
}

function fio_buffer_seq_read()
{
  echo "clean cache and disbale readahead"
  echo 1 > /proc/sys/vm/drop_caches
  disable_read_ahead ${LVM_ON}
  disable_read_ahead ${LVM_OFF}
  echo "run fio test (buffer seq read)"
  run_fio read 0
}

function fio_buffer_rand_read()
{
  echo "clean cache and disbale readahead"
  echo 1 > /proc/sys/vm/drop_caches
  disable_read_ahead ${LVM_ON}
  disable_read_ahead ${LVM_OFF}
  echo "run fio test (buffer rand read)"
  run_fio randread 0
}

function usage()
{
  echo
  echo "############################################"
  echo "iocost tests by zhaoxiaoqiang11@jd.com"
  echo "version: 1.0"
  echo "############################################"
  echo
  echo "Usage: $0 <test case>"
  echo -e "\ttest case:"
  echo -e "\tdr: direct seq read"
  echo -e "\tdw: direct seq write"
  echo -e "\tbr: buffer seq read"
  echo -e "\tbw: buffer seq write"
  echo -e "\tdrr: direct rand read"
  echo -e "\tdrw: direct rand write"
  echo -e "\tbrr: buffer rand read"
  echo -e "\tbrw: buffer rand write"
}


############################################################### main start ####################################################

if [ $# -lt 1 ]; then
  usage
  exit 1
fi

r=$(mount | grep 'on /export/iocost_' -c)
if [ $r -ne 2 ]; then
  echo "mount lvm"
  mount ${LVM_ON} ${DIRON}
  mount ${LVM_OFF} ${DIROFF}
fi

echo "set blkio qos"
echo "8:32 enable=1 ctrl=user rpct=95.00 rlat=5000 wpct=95.00 wlat=5000 min=10.00 max=10.00" > /sys/fs/cgroup/blkio/blkio.cost.qos

echo "setup cgroups"
if [ ! -e /sys/fs/cgroup/cpu/${BLKON} ];then
   cgcreate -g cpu:${BLKON}
fi

if [ ! -e /sys/fs/cgroup/cpu/${BLKOFF} ];then
   cgcreate -g cpu:${BLKOFF}
fi

cgset -r cpu.offline=1 ${BLKOFF}

if [ ! -e /sys/fs/cgroup/blkio/${BLKON} ];then
   cgcreate -g blkio:${BLKON}
fi

if [ ! -e /sys/fs/cgroup/blkio/${BLKOFF} ];then
   cgcreate -g blkio:${BLKOFF}
fi

if [ ! -e /sys/fs/cgroup/memory/${BLKON} ];then
   cgcreate -g memory:${BLKON}
fi

if [ ! -e /sys/fs/cgroup/memory/${BLKOFF} ];then
   cgcreate -g memory:${BLKOFF}
fi

cgset -r memory.limit_in_bytes=$((1024 * 1024 * 1024 * 12)) ${BLKON}
cgset -r memory.limit_in_bytes=$((1024 * 1024 * 1024 * 12)) ${BLKOFF}

echo "set blkio cost weight"
echo "set offline weight to ${WEIGHT_OFFLINE}"
cgset -r blkio.cost.weight=${WEIGHT_OFFLINE} ${BLKOFF}
echo "set online weight to ${WEIGHT_ONLINE}"
cgset -r blkio.cost.weight=${WEIGHT_ONLINE} ${BLKON}

case $1 in
    "dr")
        fio_direct_seq_read
	;;
    "dw")
        fio_direct_seq_write
	;;
    "br")
	fio_buffer_seq_read
	;;
    "bw")
	fio_buffer_seq_write
	;;
    "drr")
        fio_direct_rand_read
	;;
    "drw")
        fio_direct_rand_write
	;;
    "brr")
	fio_buffer_rand_read
	;;
    "brw")
	fio_buffer_rand_write
	;;
     *)
	usage
	;;
esac
